import pandas as pd
from statsmodels.tsa.statespace.sarimax import SARIMAX

df = pd.read_csv('Province.csv')

def preprocess(target):
    column_name = 'Province'
    target_value = target
    other_column_name = 'Confirmed'
    data = df[df[column_name] == target_value][other_column_name]
    model = SARIMAX(data, order=(1, 1, 1), seasonal_order=(0, 0, 0, 0))
    model_fit = model.fit(disp=False)
    # make prediction
    yhat = model_fit.predict(len(data), len(data))
    print(yhat)

    column_name = 'Province'
    target_value = target
    other_column_name = 'Recovered'
    data = df[df[column_name] == target_value][other_column_name]
    model = SARIMAX(data, order=(1, 1, 1), seasonal_order=(0, 0, 0, 0))
    model_fit = model.fit(disp=False)
    # make prediction
    yhat = model_fit.predict(len(data), len(data))
    print(yhat)

if __name__ == '__main__':
    preprocess('总计')